import numpy as np
import cv2
import os
from rknn.api import RKNN
from PIL import Image
from sklearn import preprocessing
from scipy.spatial.distance import pdist

os.environ['RKNN_DRAW_DATA_DISTRIBUTE'] = "1"

if __name__ == '__main__':
    im_file = 'img/1_001.jpg'
    BUILD_QUANT = False
    RKNN_MODEL_PATH = './model_data/mobilefacenet.rknn'
    if BUILD_QUANT:
        RKNN_MODEL_PATH = './model_data/mobilefacenet_quant.rknn'

    # Create RKNN object
    rknn = RKNN()

    NEED_BUILD_MODEL = True
    if NEED_BUILD_MODEL:
        print('--> config model')

        rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')

        print('done')
        print('--> Loading model')

        ret = rknn.load_onnx(model='./model_data/facenet_mobilenet.onnx', )
        if ret != 0:
            print('Load retinaface failed!')
            exit(ret)

        print('done')

        # Build model
        print('--> Building model')
        ret = rknn.build(do_quantization=BUILD_QUANT, dataset='./dataset.txt')
        if ret != 0:
            print('Build model failed!')
            exit(ret)
        print('done')

        if BUILD_QUANT:
            print('--> Accuracy analysis')
            rknn.accuracy_analysis(inputs='./dataset.txt', output_dir="./quant_result", target='rk3588')
            print('done')

        # Export rknn model
        print('--> Export RKNN model')
        ret = rknn.export_rknn(RKNN_MODEL_PATH)

        if ret != 0:
            print('Export rknn failed!')
            exit(ret)
        print('done')
    else:
        # Direct load rknn model
        print('Loading RKNN model')
        ret = rknn.load_rknn(RKNN_MODEL_PATH)
        if ret != 0:
            print('load rknn model failed.')
            exit(ret)
        print('done')

    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')

    # Set inputs
    img = cv2.imread(im_file)
    img = cv2.resize(img, (160, 160))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    image_1 = Image.open(im_file)
    image_1 = image_1.resize((160, 160), Image.BICUBIC)
    img = np.asarray(image_1, np.uint8)
    print(img.shape)

    # inference
    print('--> inference')
    outputs = rknn.inference(data_format='nhwc', inputs=[img])
    print('done')

    print(outputs)
    image_1 = Image.open("img/1_001.jpg")
    image_1 = image_1.resize((160, 160), Image.BICUBIC)
    img = np.asarray(image_1, np.uint8)
    outputs0 = np.array(rknn.inference(data_format='nhwc', inputs=[img])[0])

    image_1 = Image.open("img/1_002.jpg")
    image_1 = image_1.resize((160, 160), Image.BICUBIC)
    img = np.asarray(image_1, np.uint8)
    outputs1 = np.array(rknn.inference(data_format='nhwc', inputs=[img])[0])

    l1 = np.linalg.norm(outputs1 - outputs0, axis=1)
    print("l1 %f" % l1)
    cosSim = 1 - pdist(np.vstack([outputs1, outputs0]), 'cosine')
    print("pdist %f" % cosSim)
    outputs1 = preprocessing.normalize(outputs1, norm='l2')
    outputs0 = preprocessing.normalize(outputs0, norm='l2')
    l1 = np.linalg.norm(outputs1 - outputs0, axis=1)
    print("after l2 l1 %f" % l1)

    rknn.release()
